from http import HTTPStatus
from urllib.parse import urlparse, unquote
from pathlib import PurePosixPath
import requests
from dashscope import ImageSynthesis
import logging
import os
from dotenv import load_dotenv
import dashscope

# 定义日志输出格式
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s'
)

# (1)加载 .env 文件中的环境变量
dotenv_path = os.path.join(os.path.dirname(__file__), '../.env')
load_dotenv(dotenv_path=dotenv_path)


def sample_block_call(input_prompt,dst_path):
    # (1-2)从环境变量中获取 DASHSCOPE_API_KEY
    api_key = os.getenv("DASHSCOPE_API_KEY") # 从环境变量中获取 DASHSCOPE_API_KEY
    print(api_key)
    #  (1-3)确保api_key被正确设置
    if not api_key:
        raise ValueError("DASHSCOPE_API_KEY is not set in the environment variables.")

    # (1-4)设置dashscope的api_key
    dashscope.api_key = api_key

    # (2) 目录校验：获取dst_path的目录部分，并将其赋值给变量dst_dir
    # 检查dst_dir是否存在，如果不存在，使用os.makedirs函数创建该目录及其所有父目录
    if not dst_path.endswith('.png'):
        raise ValueError("Only support .png file")
    dst_dir = os.path.dirname(dst_path)
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
    # （3）传入模型名称、文本提示、图像大小和生成步骤数等参数，发起图像生成请求，并将响应结果存储在变量rsp中
    rsp = ImageSynthesis.call(model=model,
                              prompt=input_prompt,
                              size='1024*1024')
    if rsp.status_code == HTTPStatus.OK:
        print(rsp.output)
        print(rsp.usage)
        # 两行代码打印生成的图像输出信息和资源使用情况
        for result in rsp.output.results:
            #以二进制写模式打开dst_path指定的文件，准备写入生成的图像数据。
            file_name = PurePosixPath(unquote(urlparse(result.url).path)).parts[-1]
            with open('./%s' % file_name, 'wb+') as f:
                # 从结果项的URL下载图像数据，并将其写入打开的文件中
                f.write(requests.get(result.url).content)
    else:
        print('Failed, status_code: %s, code: %s, message: %s' %
              (rsp.status_code, rsp.code, rsp.message))


def sample_async_call(input_prompt):
    rsp = ImageSynthesis.async_call(model=model,
                                    prompt=input_prompt,
                                    size='1024*1024')
    if rsp.status_code == HTTPStatus.OK:
        print(rsp.output)
        print(rsp.usage)
    else:
        print('Failed, status_code: %s, code: %s, message: %s' %
              (rsp.status_code, rsp.code, rsp.message))
    status = ImageSynthesis.fetch(rsp)
    if status.status_code == HTTPStatus.OK:
        print(status.output.task_status)
    else:
        print('Failed, status_code: %s, code: %s, message: %s' %
              (status.status_code, status.code, status.message))

    rsp = ImageSynthesis.wait(rsp)
    if rsp.status_code == HTTPStatus.OK:
        print(rsp.output)
    else:
        print('Failed, status_code: %s, code: %s, message: %s' %
              (rsp.status_code, rsp.code, rsp.message))

# dashscope是一个提供AI服务的库，而ImageSynthesis模块可能是用于图像合成或生成的功能
# 该函数用于根据输入的文本提示生成图像，并将生成的图像保存到指定路径
if __name__ == '__main__':
    model = "flux-dev"  # 指定模型
    prompt = "小女孩在草丛中玩耍"
    sample_block_call(prompt,"./tmp/output.png")
    # sample_async_call(prompt)